#!/usr/bin/env python3

import json
import yaml
import sys
import os
from pathlib import Path
import numpy as np
import cv2
from typing import Dict
from copy import deepcopy
sys.path.extend([os.path.abspath(Path(__file__).parent.parent), os.path.abspath(Path(__file__).parent.parent / "RoboTwin"), os.path.abspath(Path(__file__).parent.parent / "inject_pcd" / "adjust")])
sys.path.append(os.path.abspath(Path(__file__).parent.parent / "inject_pcd" / "adjust"))
sys.path.append(os.path.abspath(Path(__file__).parent.parent / "inject_pcd"))

from intrinsics import intrinsics_real,intrinsics_sim
from envs.utils.pcd_trans import pcd_transform, pcd2ndarray, depth_to_point_cloud, ndarray2pcd
from debugger.save_file import save_info
import open3d as o3d
from cam_yaml_new import get_pose_from_yaml, create_transformation_matrix_from_plf
from envs.utils.vis_pcd import catpcd
from fpsPcd import fpsPcd
def injPcd_test(real_rgbd_path, mbr_deploy_rgbd_path, transform_yaml_path, sim_pcd_path):
    real_rgbd = np.load(real_rgbd_path)
    mbr_deploy_rgbd = np.load(mbr_deploy_rgbd_path)
    mbr_sim_mask = mbr_deploy_rgbd[:, :, -1]
    mbr_sim_mask = cv2.resize(mbr_sim_mask, (real_rgbd.shape[1], real_rgbd.shape[0]), interpolation=cv2.INTER_NEAREST)
    assert real_rgbd[:,:,-1].shape == mbr_sim_mask.shape, f"real_rgbd.shape != mbr_sim_mask.shape, {real_rgbd.shape}[:,:,-1] != {mbr_sim_mask.shape}"

    real_rgbd = real_rgbd.copy()
    print("real_rgbd", real_rgbd.shape)
    save_info("depth", "depth_masked", real_rgbd[:,:,-1].copy(),mbr_sim_mask)
    
    print("real_rgbd masked", real_rgbd.shape)
    print("mask", mbr_sim_mask.shape)
   
    # 生成点云并带颜色
    object_pcd_mbrframe = depth_to_point_cloud(real_rgbd, intrinsics_real)
    save_info("pcd", "object_pcd_mbrframe", pcd2ndarray(object_pcd_mbrframe))

    with open(transform_yaml_path, 'r', encoding='utf-8') as f:
        yaml_data = yaml.safe_load(f)
    pose = get_pose_from_yaml(yaml_data, cam_name='midBack_real_camera')
    p = pose['position']
    l = pose['left']
    f = pose['forward']
    trans_mat = create_transformation_matrix_from_plf(p, l, f)
    print("trans_mat", trans_mat)
    obj_pcd = pcd_transform(trans_mat, object_pcd_mbrframe,True)

    sim_pcd = o3d.io.read_point_cloud(sim_pcd_path)#sim pcd path is a camera frame pcd
    pose_sim = get_pose_from_yaml(yaml_data, cam_name='midBack_camera')
    p_sim = pose_sim['position']
    l_sim = pose_sim['left']
    f_sim = pose_sim['forward']
    trans_mat_sim = create_transformation_matrix_from_plf(p_sim, l_sim, f_sim)
    sim_pcd = pcd_transform(trans_mat_sim, sim_pcd,True)
    result_pcd = catpcd(obj_pcd, sim_pcd)
    result_pcd = fpsPcd({"repointcloud":result_pcd}, 4096)["repointcloud"]
    return result_pcd
def rgbd_pack_ud(unified_data:Dict,cam_key='real_camera'):
    rgb = unified_data['cameras'][cam_key]['rgb']
    depth = unified_data['cameras'][cam_key]['depth']
    return rgbd_pack(rgb,depth)
def rgbd_pack(rgb,depth):
    depth_extended = np.expand_dims(depth, axis=-1)
    rgbd = np.concatenate([rgb, depth_extended], axis=-1)
    return rgbd
def rgbd_unpack(rgbd):
    rgb = rgbd[:,:,:3]
    depth = rgbd[:,:,3]
    return rgb,depth
def intrin_mat_to_tuple(intrin_mat):
    return (intrin_mat[0,0],intrin_mat[1,1],intrin_mat[0,2],intrin_mat[1,2])
def add_real_camera_cfg(unified_data:Dict):
    unified_data['cameras']['real_camera']['intrinsic_cv'] = deepcopy(unified_data['cameras']['midBack_real_camera']['intrinsic_cv'])
    unified_data['cameras']['real_camera']['intrinsic_cv'][0:2,:] =2*unified_data['cameras']['midBack_real_camera']['intrinsic_cv'][0:2,:]
    unified_data['cameras']['real_camera']['cam2world_gl'] = deepcopy(unified_data['cameras']['midBack_real_camera']['cam2world_gl'])
    if False:
        print("midBack_real_camera", unified_data['cameras']['midBack_real_camera']['intrinsic_cv'],"intrinsics_sim", intrinsics_sim)
        print("real_camera", unified_data['cameras']['real_camera']['intrinsic_cv'],"intrinsics_real", intrinsics_real)
        in_real_mat = np.array(unified_data['cameras']['real_camera']['intrinsic_cv'])
        in_real = intrin_mat_to_tuple(in_real_mat)
        print("intrinsics_real again", in_real)

def add_real_camera(unified_data:Dict,real_rgb,real_depth):
    unified_data['cameras']['real_camera'] = deepcopy(unified_data['cameras']['midBack_real_camera'])
    unified_data['cameras']['real_camera']['rgb'] = real_rgb
    unified_data['cameras']['real_camera']['depth'] = real_depth#NOTE THAT real camera rgb depth shape are 480*640 instead of 240*320
CNT=0
def mask_real_arm_ud(unified_data:Dict):
    global CNT
    real_rgbd = rgbd_pack_ud(unified_data)
    mask = np.array(unified_data['mbr_deploy_mask'])
    mask = mask.astype(np.uint8)
    mask = cv2.resize(mask, (real_rgbd.shape[1], real_rgbd.shape[0]), interpolation=cv2.INTER_NEAREST)
    assert real_rgbd[:,:,-1].shape == mask.shape, f"real_rgbd.shape != mbr_sim_mask.shape, {real_rgbd.shape}[:,:,-1] != {mbr_sim_mask.shape}"
    real_rgbd = real_rgbd.copy()
    # real_rgbd[mask.astype(bool),-1] = 0
    save_info("depth", Path(__file__).parent.parent / "zarrtraj"/f"depth_masked_{CNT}", real_rgbd[:,:,-1].copy(),mask)
    CNT+=1
    return real_rgbd
def get_intrin_ud(unified_data:Dict,cam_key='real_camera'):
    return intrin_mat_to_tuple(unified_data['cameras']['real_camera']['intrinsic_cv'])
def get_cam2world_o3d_ud(unified_data:Dict,cam_key='real_camera'):
    mat=np.array(unified_data['cameras'][cam_key]['cam2world_gl'])
    mat[:,1:3] = -mat[:,1:3]
    return mat
def injPcd(unified_data:Dict) -> np.ndarray:
    """
    called by test_shoe_place_task.py
    params:
        unified_data: dict containing real_rgbd, mbr_deploy_mask, cameras, repointcloud
        - real_rgbd: np.ndarray, shape (H,W,4)
        - mbr_deploy_mask: np.ndarray, shape (H,W)
        - repointcloud: np.ndarray, shape (N,6)
    return:
        np.ndarray: combined point cloud
    """
    if unified_data['cameras'].get('real_camera') is None or unified_data['cameras'].get('real_camera')['rgb'] is None or unified_data['cameras'].get('real_camera')['depth'] is None:
        rgbd = np.load("/home/algo/geyiheng/inject_pcd/adjust/test1/real.npy")
        real_rgb,real_depth = rgbd_unpack(rgbd)
        add_real_camera(unified_data,real_rgb,real_depth)
    add_real_camera_cfg(unified_data)
    obj_rgbd = mask_real_arm_ud(unified_data)
    obj_pcd_local = depth_to_point_cloud(obj_rgbd, get_intrin_ud(unified_data),True)
    # save_info("pcd", "object_pcd_mbrframe", obj_pcd_local)
    obj_pcd_world = pcd_transform(get_cam2world_o3d_ud(unified_data), obj_pcd_local,True)
    unified_data['injpointcloud'] = catpcd(obj_pcd_world, unified_data['repointcloud'])
    return unified_data

if __name__ == "__main__":
    real_rgbd_path = "/home/algo/geyiheng/inject_pcd/adjust/test1/real.npy"
    sim_pcd_path = "/home/algo/geyiheng/inject_pcd/adjust/test1/ref.pcd"
    mbr_deploy_rgbd_path = "/home/algo/geyiheng/inject_pcd/adjust/mbr_deploy_mask/midBack_real_camera_rgbd.npy"
    transform_json_path = "/home/algo/geyiheng/inject_pcd/adjust/test1/to_adjust.json"
    transform_yaml_path = '/home/algo/geyiheng/inject_pcd/adjust/test1/cam_pfl.yaml'

    result_pcd = injPcd_test(real_rgbd_path, mbr_deploy_rgbd_path, transform_yaml_path,sim_pcd_path)
    save_info("pcd", "processed", result_pcd)
    print("处理完成，结果点云形状:", result_pcd.shape)
